import time
import torch
import random

def match(key, seq):
    i = 0
    for k, w in zip(key, seq):
        if k != w:
            break
        i += 1
    return i

def match2(t1: torch.Tensor, t2: torch.Tensor) -> int:
    # Ensure same shape for comparison: flatten and get min length
    t1_flat = t1.flatten()
    t2_flat = t2.flatten()
    min_len = min(t1.size(0), t2.size(0))

    # Compare elements and find first mismatch
    diff = t1[:min_len] != t2[:min_len]
    # print(diff)
    mismatch_indices = torch.nonzero(diff)
    # print(mismatch_indices)

    if mismatch_indices.numel() == 0:
        return min_len  # All matched up to min_len
    else:
        return mismatch_indices[0].item()

def generate_random_tensors(min_len=3, max_len=10, min_val=0, max_val=100):
    len1 = random.randint(min_len, max_len)
    len2 = random.randint(min_len, max_len)

    t1 = torch.randint(min_val, max_val, (len1,))
    t2 = torch.randint(min_val, max_val, (len2,))

    return t1, t2

start = time.time()

t1, t2 = 0, 0

for i in range(100):
    key, seq = generate_random_tensors(min_len=1000, max_len=10000)
    start = time.time()
    p1 = match2(key, seq)
    t1 += time.time() - start
    start = time.time()
    p2 = match(key, seq)
    t2 += time.time() - start
    assert p1 == p2, f"Mismatch: {p1} != {p2}"


print("Total time for 100 iterations:", t1, t2)